# Ultralytics YOLO 🚀, AGPL-3.0 license
"""Convolution modules."""

import math

import numpy as np
import torch
import torch.nn as nn

__all__ = (
    "MobileViTBackbone3",
)

import torch
import torch.nn as nn
import timm

# 可见光网络
class VisibleNet(nn.Module):
    def __init__(self, step=1, *args, **kwargs):
        super().__init__(*args, **kwargs)

        mobilevit = timm.create_model('mobilevit_s', pretrained=False)

        # 第一层
        self.layer1_1 = list(mobilevit.children())[0]
        self.layer1_2 = list(mobilevit.children())[1][0:3]
        # original_conv1 = self.layer1_2[0]  # 获取第一层卷积


        # 第二层
        self.layer2_1 = list(mobilevit.children())[1][3]

        # 第三层
        self.layer3_1 = list(mobilevit.children())[1][4:]
        self.layer3_2 = list(mobilevit.children())[2]

        # 以640 * 640 输入为例，输出结果为80*80
        if step == 1:
            self.layer = nn.Sequential(
                self.layer1_1,
                *self.layer1_2,
            )

        # 输出为 40*40
        elif step == 2:
            self.layer = self.layer2_1

        # 输出为 20*20
        elif step == 3:
            self.layer = nn.Sequential(
                *self.layer3_1, self.layer3_2
            )

    def forward(self, x):
        # print(f"VisibleNet 输入形状为：{x.shape}")
        return self.layer(x)


class MobileViTBackbone3(nn.Module):
    def __init__(self, step=1, layer1_r=1, layer2_r=1):
        super(MobileViTBackbone3, self).__init__()

        self.layer1_rate = layer1_r #两个通道的比例
        self.layer2_rate = layer2_r

        # 分别加载这两部分的网络
        self.visibleNet = VisibleNet(step=step)

        pass

    def forward(self, x):
        visible_x_out = self.visibleNet(x)
        return visible_x_out